import requests
from bs4 import BeautifulSoup
import json
import re
import mysql.connector


def get_platforms(id):
    url = f'https://attack.mitre.org/techniques/{id}/'
    try:
        response = requests.get(url, proxies={"http": "http://127.0.0.1:7890", "https": "http://127.0.0.1:7890"})
    except:
        print('Error: ', url)
        return 'website error'
    technique_soup = BeautifulSoup(response.content, 'html.parser')
    # 提取平台信息
    platforms_span = technique_soup.find('span', class_='h5 card-title', text='Platforms:\xa0')
    if platforms_span:
        platforms_div = platforms_span.find_parent('div')
        platforms = platforms_div.get_text(strip=True).replace('Platforms:', '').strip()
    else:
        platforms = 'No platforms available'
    return platforms

def get_tacticDescription_techniques(soup):
    # 提取描述信息并去掉[]包围的数字
    tactic_description_div = soup.find('div', class_='description-body')
    if tactic_description_div:
        tactic_description_text = tactic_description_div.get_text()
        tactic_description = re.sub(r'\[\d+\]', '', tactic_description_text).strip()
    else:
        tactic_description = 'No description available'
    # 获取相关技术
    technique_list = []
    subtechnique_list = []
    # 遍历表格的每一行
    technique_table = soup.find('table', class_='table-techniques')
    for row in technique_table.find('tbody').find_all('tr'):
        if row.get('class') == ['technique']:
            # 如果subtechnique_list不为空，说明上一个技术有子技术，将子技术列表添加到上一个技术中
            if subtechnique_list != []:
                technique_list[-1]['subtechniques'] = subtechnique_list
            # 提取ID、Name和Description
            technique_id = row.find_all('td')[0].get_text(strip=True)
            technique_name = row.find_all('td')[1].get_text(strip=True)
            technique_description = row.find_all('td')[2].get_text(strip=True)
            technique_platform = get_platforms(technique_id)
            technique_list.append({
                'technique_mitre_id': technique_id,
                'technique_name': technique_name,
                'technique_platform': technique_platform,
                'technique_description': technique_description,
            })
            # 清空子技术列表，准备存储下一个技术的子技术
            subtechnique_list = []
            continue
        
        if 'sub' in row.get('class', []) and 'technique' in row.get('class', []):
            # 提取ID、Name和Description
            subtechnique_id_temp = row.find_all('td')[1].get_text(strip=True)
            # 提取子技术platform
            subtechnique_platform = get_platforms(technique_id + '/' + subtechnique_id_temp.lstrip('.'))
            print(subtechnique_platform)
            # 拼接出子技术ID
            subtechnique_id = technique_id + subtechnique_id_temp
            subtechnique_name = row.find_all('td')[2].get_text(strip=True)
            subtechnique_description = row.find_all('td')[3].get_text(strip=True)
            subtechnique_list.append({
                'subtechnique_mitre_id': subtechnique_id,
                'subtechnique_name': subtechnique_name,
                'subtechnique_description': subtechnique_description,
                'subtechnique_platform': subtechnique_platform
            })
        
    return tactic_description, technique_list

def get_attck_list():
    # 发送GET请求并获取网页内容
    url = 'https://attack.mitre.org/tactics/enterprise/'
    response = requests.get(url, proxies={"http": "http://127.0.0.1:7890", "https": "http://127.0.0.1:7890"})
    # 使用BeautifulSoup解析网页内容
    soup = BeautifulSoup(response.content, 'html.parser')
    # 所有战术
    tactic_table = soup.find('table', class_='table table-bordered table-alternate mt-2')
    # 创建一个空战术数组用于存储数据
    attck_list =[]
    # 遍历表格的每一行
    for row in tactic_table.find('tbody').find_all('tr'):
        # 提取ID和Name
        tactic_id = row.find_all('td')[0].get_text(strip=True)
        tactic_name = row.find_all('td')[1].get_text(strip=True)
        # 根据ID获取战术描述和相关技术
        url = f'https://attack.mitre.org/tactics/{tactic_id}/'
        response = requests.get(url, proxies={"http": "http://127.0.0.1:7890", "https": "http://127.0.0.1:7890"})
        soup = BeautifulSoup(response.content, 'html.parser')
        # 获取战术描述和技术（技术里有子技术，如果不需要，请注释掉get_tacticDescription_techniques函数中的子技术部分）
        tactic_description, techniques= get_tacticDescription_techniques(soup)
        attck_list.append({'tactic_mitre_id': tactic_id, 'tactic_name': tactic_name, 'tactic_description': tactic_description, 'techniques': techniques})
    
    return attck_list

# 将ATT&CK插入数据库 这里有两个表，attck_tactic、attck_technique
def insert_database(attck_list):
    # # 从JSON文件中读取数据
    # with open('attck_list.json', 'r', encoding='utf-8') as f:
    #     attck_list = json.load(f)
    # 连接到MySQL数据库
    conn = mysql.connector.connect(
        host='localhost',  # 替换为你的MySQL主机
        user='root',  # 替换为你的MySQL用户名
        password='123456',  # 替换为你的MySQL密码
        database='mydatabase'  # 替换为你的MySQL数据库名
    )
    cursor = conn.cursor()

    # 插入tactic数据
    for tactic in attck_list:
        cursor.execute('''
        INSERT INTO attck_tactic (tactic_mitre_id, tactic_name, tactic_description)
        VALUES (%s, %s, %s)
        ''', (tactic['tactic_mitre_id'], tactic['tactic_name'], tactic['tactic_description']))
        tactic_id = cursor.lastrowid  # 获取插入的tactic的自增ID

        # 插入technique数据
        for technique in tactic['techniques']:
            cursor.execute('''
            INSERT INTO attck_technique (technique_mitre_id, technique_name, technique_platform, technique_description, fk_technique_tactic_id)
            VALUES (%s, %s, %s, %s, %s)
            ''', (technique['technique_mitre_id'], technique['technique_name'], technique['technique_platform'], technique['technique_description'], tactic_id))

    # 提交事务
    conn.commit()

    # 关闭数据库连接
    conn.close()

# 获取ATT&CK并写入json文件
attck_list = get_attck_list()
with open('attck_result.json', 'w', encoding='utf-8') as f:
    json.dump(attck_list, f, ensure_ascii=False, indent=4)
# 这里是存库的，不需要请注释掉
# insert_database(attck_list)




# 也可以直接从json中读取然后存库
# attck_list = []
# with open('attck_result.json', 'r', encoding='utf-8') as f:
#     attck_list = json.load(f)
# insert_database(attck_list)
